/**
 * @author: mattwang@tencent.com
 * @date: 2012-10-10
 */

#include <ostream>
#include <iostream>
#include <fstream>
#include <math.h>

#include <map>
#include <string>
#include	 <type_def.h>
#include <param_def.h>

#include "impl/document.h"
#include "impl/kernel.h"
#include "impl/pd_ipm_parm.h"
#include "impl/pd_ipm.h"
#include "impl/parallel_interface.h"
#include "impl/matrix.h"
#include "impl/matrix_manipulation.h"
#include "impl/util.h"
#include "timer.h"
#include <svm_core.h>

// Incomplete Cholesky Fatorization related options
double FLAGS_fact_threshold = 1.0e-5;
double FLAGS_rank_ratio = 1.0;

// Kernel function related options
int FLAGS_kernel_type = 2;
double FLAGS_gamma = 1.0;
int FLAGS_poly_degree = 3;
double FLAGS_poly_coef = 1.0;
double FLAGS_poly_const = 1.0;

// Learning related options
double FLAGS_zero_threshold = 1.0e-9;
double FLAGS_sv_threshold = 1.0e-4;
double FLAGS_hyper_parm = 1.0;
double FLAGS_positive_weight = 1.0;
double FLAGS_negative_weight = 1.0;

// IPM related options
double FLAGS_feasible_threshold = 1.0e-3;
double FLAGS_surrogate_gap_threshold = 1.0e-3;
int FLAGS_max_iteration = 200;
double FLAGS_mu_factor = 10.0;

// Where to save the result
string FLAGS_model_path = ".";

// Miscs
bool FLAGS_verbose = false;
double FLAGS_save_interval = 600.0;
bool FLAGS_failsafe = false;

string FLAGS_output_path = ".";
int FLAGS_batch_size = 10000;

int SvmCore::train(string filename) {
	Kernel kernel;
	Model model;
	Document doc;
	PrimalDualIPMParameter ipm_parameter;

	// Sets the commandline options
	ipm_parameter.threshold = FLAGS_fact_threshold;
	ipm_parameter.rank_ratio = FLAGS_rank_ratio;
	kernel.set_kernel_type(static_cast<Kernel::KernelType>(FLAGS_kernel_type));
	kernel.set_rbf_gamma(FLAGS_gamma);
	kernel.set_poly_degree(FLAGS_poly_degree);
	kernel.set_coef_lin(FLAGS_poly_coef);
	kernel.set_coef_const(FLAGS_poly_const);
	ipm_parameter.epsilon_x = FLAGS_zero_threshold;
	ipm_parameter.epsilon_sv = FLAGS_sv_threshold;
	ipm_parameter.hyper_parm = FLAGS_hyper_parm;
	strncpy(ipm_parameter.model_path, FLAGS_model_path.c_str(), sizeof(ipm_parameter.model_path));
	ipm_parameter.weight_positive = FLAGS_positive_weight;
	ipm_parameter.weight_negative = FLAGS_negative_weight;
	ipm_parameter.feas_thresh = FLAGS_feasible_threshold;
	ipm_parameter.sgap = FLAGS_surrogate_gap_threshold;
	ipm_parameter.max_iter = FLAGS_max_iteration;
	ipm_parameter.mu_factor = FLAGS_mu_factor;
	ipm_parameter.verb = FLAGS_verbose;
	ipm_parameter.save_interval = FLAGS_save_interval;
	ipm_parameter.tradeoff = 0;

	ParallelInterface *mpi = ParallelInterface::GetParallelInterface();
	ParallelMatrix icf_result;
	int rank = static_cast<int>(doc.GetGlobalNumberRows() * ipm_parameter.rank_ratio);
	if (FLAGS_failsafe) {
		int local_succeed = 0;
		int global_succeed = 0;
		if (icf_result.Load(ipm_parameter.model_path, "ICF"))
			local_succeed = 1;
		mpi->AllReduce(&local_succeed, &global_succeed, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);
		if (global_succeed != mpi->GetNumProcs()) {
			// Performs Incomplete Cholesky Fatorization to
			// reduce the matrix dimension.
			if (mpi->GetProcId() == 0)
				cout << "Performing ICF ... \t" << endl;
			Timer t1;
			t1.Start();
			MatrixManipulation::ICF(doc, kernel, doc.GetGlobalNumberRows(), rank, ipm_parameter.threshold, &icf_result);
			t1.Stop();
			if (mpi->GetProcId() == 0)
				cout << "ICF completed in " << t1.total() << " seconds" << endl;
			if (mpi->GetProcId() == 0)
				cout << "Saving H ..." << endl;
			Timer t2;
			t2.Start();
			icf_result.Save(ipm_parameter.model_path, "ICF");
			t2.Stop();
			if (mpi->GetProcId() == 0) {
				cout << "H saved in " << t2.total() << " seconds" << endl;
			}
		}
	} else {
		if (mpi->GetProcId() == 0)
			cout << "Performing ICF ... " << endl;
		Timer t1;
		t1.Start();
		MatrixManipulation::ICF(doc, kernel, doc.GetGlobalNumberRows(), rank, ipm_parameter.threshold, &icf_result);
		t1.Stop();
		if (mpi->GetProcId() == 0) {
			cout << "ICF completed in " << t1.total() << " seconds" << endl;
		}

	}
	PrimalDualIPM ipm_solver;
	ipm_solver.Solve(ipm_parameter, icf_result, doc, &model_, FLAGS_failsafe);
	icf_result.Destroy();

	model_.set_kernel(kernel);
	model_.ComputeB(ipm_parameter);
	model_.Save(ipm_parameter.model_path, "model");
	mpi->Barrier(MPI_COMM_WORLD);

	return 0;
}

int SvmCore::predict(string test_file) {
	int chunk_size = FLAGS_batch_size;
	string predict_filename = FLAGS_output_path + "/PredictResult";
	ParallelInterface* interface = ParallelInterface::GetParallelInterface();
	int myid = interface->GetProcId();

	const Kernel *kernel = model_.kernel();
	const SupportVector *support_vector = model_.support_vector();

	double *local_prediction = new double[chunk_size];
	int *label = new int[chunk_size];
	double *global_prediction = NULL;
	if (myid == 0)
		global_prediction = new double[chunk_size];

	int num_total_document = 0;
	int num_positive_positive = 0;  // Number of samples whose original class
	// label is positive and predicted as
	// positive.
	int num_positive_negative = 0;
	int num_negative_positive = 0;
	int num_negative_negative = 0;
	int num_parsed_samples = 0;

	string line;

	File *outputbuffer_predict = NULL;

	if (myid == 0) {
		// Open the predict file for writing.
		outputbuffer_predict = File::Open(predict_filename.c_str(), "w");
	}

	File* reader = File::OpenOrDie(test_file.c_str(), "r");
	bool end_of_file = false;
	char sz_line[256];
	while (!end_of_file) {
		// Read a line
		end_of_file = !reader->ReadLine(&line);

		// If the local prediction of chunk_size samples is ready or if there are no
		// samples left to predict, performs a reduce to obtain global prediction
		// result.
		if (num_parsed_samples == chunk_size || end_of_file) {
			// Processor 0 will have the global prediction result.
			interface->Reduce(local_prediction, global_prediction, num_parsed_samples, MPI_DOUBLE, MPI_SUM, 0,
					MPI_COMM_WORLD);
			if (myid == 0) {
				for (int i = 0; i < num_parsed_samples; i++) {
					double predicted_label = global_prediction[i] + support_vector->b;

					// Output predicted values.
					size_t b = snprintf(sz_line, sizeof(sz_line), "%d %f\n", label[i], predicted_label);
					CHECK(outputbuffer_predict->Write(sz_line, b) == b);

					// Updates the counters
					if (label[i] > 0) {
						if (predicted_label > 0)
							++num_positive_positive;
						else
							++num_positive_negative;
					} else {
						if (predicted_label > 0)
							++num_negative_positive;
						else
							++num_negative_negative;
					}
				}

				// Flush buffer per batch
				if (!outputbuffer_predict->Flush()) {
					cerr << "Flushing " << predict_filename << " failed!";
					exit(1);
				}

				// Print the number of predicted samples.
				cout << "Finished: " << num_total_document << endl;
			}
			num_parsed_samples = 0;

		}
		if (end_of_file)
			break;

		// Parses one sample
		const char* start = line.c_str();
		if (!SplitOneIntToken(&start, " ", &label[num_parsed_samples])) {
			return -1;
		}
		vector<pair<string, string> > kv_pairs;
		SplitStringIntoKeyValuePairs(string(start), ":", " ", &kv_pairs);
		Sample sample;
		sample.label = label[num_parsed_samples];
		sample.two_norm_sq = 0.0;
		vector<pair<string, string> >::const_iterator pair_iter;
		for (pair_iter = kv_pairs.begin(); pair_iter != kv_pairs.end(); ++pair_iter) {
			Feature feature;
			feature.id = atoi(pair_iter->first.c_str());
			feature.weight = atof(pair_iter->second.c_str());
			sample.features.push_back(feature);
			sample.two_norm_sq += feature.weight * feature.weight;
		}

		// Tests the sample based on local support vectors
		double predicted_label = 0.0;
		double temp;
		for (int j = 0; j < support_vector->num_sv; ++j) {
			temp = kernel->CalcKernel(sample, support_vector->sv_data_test[j]);
			predicted_label += support_vector->sv_alpha[j] * temp;
		}
		local_prediction[num_parsed_samples] = predicted_label;

		// Moves on to next sample
		++num_parsed_samples;
		++num_total_document;
	}

	// Cleans up
	reader->Close();
	delete reader;
	delete[] local_prediction;
	delete[] label;
	if (myid == 0) {
		delete[] global_prediction;
		CHECK(outputbuffer_predict->Close());
		delete outputbuffer_predict;
	}

	// Computes some statistics
	// Record the numbers
	// Calculate the precision/recall and accuracy
	int correct = num_positive_positive + num_negative_negative;
	int incorrect = num_negative_positive + num_positive_negative;

	// Prints the statistical information.
	interface->Barrier(MPI_COMM_WORLD);

	return 0;
}

